{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "67ce3338-182c-4d0d-815f-8eec476824d5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import sys\n",
    "sys.path.append('./utils')\n",
    "sys.path.append('./utils/APIs')\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import timm\n",
    "import argparse\n",
    "from Config import config\n",
    "from utils.common import read_from_file, save_model, write_to_file, train_val_split\n",
    "from utils.DataProcess import Processor\n",
    "from Trainer import Trainer\n",
    "from PreTrainer import PreTrainer\n",
    "from Models.mamba.models_insect import build_vssm_models_ as mamba_model\n",
    "from Models.OTEModel import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7f93789f-6168-4e81-ac93-add5171d5331",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TextModel: huawei-noah/TinyBERT_General_4L_312D, ImageModel: WAVM, FuseModel: CMAT\n"
     ]
    }
   ],
   "source": [
    "# args\n",
    "# 训练模型\n",
    "do_train = True\n",
    "# 预测测试集数据\n",
    "do_test = True\n",
    "# 已经训练好的模型路径\n",
    "load_model_path = './save_models/CMAT/pytorch_model.bin'\n",
    "# 设置学习率（后面会通过GRW-DSA寻找最优学习率）\n",
    "lr = 3e-6\n",
    "# 设置权重衰减\n",
    "weight_decay = 1e-2\n",
    "# 设置训练轮数\n",
    "epoch = 20\n",
    "# 文本分析模型\n",
    "text_pretrained_model = 'huawei-noah/TinyBERT_General_4L_312D'\n",
    "cv_pretrained_model = 'WAVM'\n",
    "# 融合模型类别\n",
    "fuse_model_type = 'CMAT'\n",
    "# 仅用文本预测\n",
    "text_only = False\n",
    "# 仅用图像预测\n",
    "img_only = False\n",
    "\n",
    "config.learning_rate = lr\n",
    "config.weight_decay = weight_decay\n",
    "config.epoch = epoch\n",
    "config.bert_name = text_pretrained_model\n",
    "config.resnet_name = cv_pretrained_model\n",
    "config.fuse_model_type = fuse_model_type\n",
    "config.load_model_path = load_model_path\n",
    "config.only = 'img' if img_only else None\n",
    "config.only = 'text' if text_only else None\n",
    "if img_only and text_only: config.only = None\n",
    "print('TextModel: {}, ImageModel: {}, FuseModel: {}'.format(config.bert_name,config.resnet_name, config.fuse_model_type))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5792737a-6e11-46e5-b4ef-d74f2b9d7361",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-09 20:25:29.052923: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-03-09 20:25:29.109318: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-03-09 20:25:30.003112: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully load ckpt ./vssm_tiny_0230_ckpt_epoch_262.pth\n",
      "_IncompatibleKeys(missing_keys=['layers.0.blocks.0.mutli_head_atten.in_proj_weight', 'layers.0.blocks.0.mutli_head_atten.in_proj_bias', 'layers.0.blocks.0.mutli_head_atten.out_proj.weight', 'layers.0.blocks.0.mutli_head_atten.out_proj.bias', 'layers.0.blocks.0.conv33conv33conv11.0.weight', 'layers.0.blocks.0.conv33conv33conv11.0.bias', 'layers.0.blocks.0.conv33conv33conv11.1.weight', 'layers.0.blocks.0.conv33conv33conv11.1.bias', 'layers.0.blocks.0.conv33conv33conv11.1.running_mean', 'layers.0.blocks.0.conv33conv33conv11.1.running_var', 'layers.0.blocks.0.conv33conv33conv11.3.weight', 'layers.0.blocks.0.conv33conv33conv11.3.bias', 'layers.0.blocks.0.conv33conv33conv11.4.weight', 'layers.0.blocks.0.conv33conv33conv11.4.bias', 'layers.0.blocks.0.conv33conv33conv11.4.running_mean', 'layers.0.blocks.0.conv33conv33conv11.4.running_var', 'layers.0.blocks.0.conv33conv33conv11.6.weight', 'layers.0.blocks.0.conv33conv33conv11.6.bias', 'layers.0.blocks.0.sk_conv.conv.0.weight', 'layers.0.blocks.0.sk_conv.conv.0.bias', 'layers.0.blocks.0.sk_conv.conv.1.weight', 'layers.0.blocks.0.sk_conv.conv.1.bias', 'layers.0.blocks.0.sk_conv.conv.2.weight', 'layers.0.blocks.0.sk_conv.conv.2.bias', 'layers.0.blocks.0.sk_conv.conv.3.weight', 'layers.0.blocks.0.sk_conv.conv.3.bias', 'layers.0.blocks.0.sk_conv.fc.0.weight', 'layers.0.blocks.0.sk_conv.fc.0.bias', 'layers.0.blocks.0.sk_conv.fc.2.weight', 'layers.0.blocks.0.sk_conv.fc.2.bias', 'layers.0.blocks.1.mutli_head_atten.in_proj_weight', 'layers.0.blocks.1.mutli_head_atten.in_proj_bias', 'layers.0.blocks.1.mutli_head_atten.out_proj.weight', 'layers.0.blocks.1.mutli_head_atten.out_proj.bias', 'layers.0.blocks.1.conv33conv33conv11.0.weight', 'layers.0.blocks.1.conv33conv33conv11.0.bias', 'layers.0.blocks.1.conv33conv33conv11.1.weight', 'layers.0.blocks.1.conv33conv33conv11.1.bias', 'layers.0.blocks.1.conv33conv33conv11.1.running_mean', 'layers.0.blocks.1.conv33conv33conv11.1.running_var', 'layers.0.blocks.1.conv33conv33conv11.3.weight', 'layers.0.blocks.1.conv33conv33conv11.3.bias', 'layers.0.blocks.1.conv33conv33conv11.4.weight', 'layers.0.blocks.1.conv33conv33conv11.4.bias', 'layers.0.blocks.1.conv33conv33conv11.4.running_mean', 'layers.0.blocks.1.conv33conv33conv11.4.running_var', 'layers.0.blocks.1.conv33conv33conv11.6.weight', 'layers.0.blocks.1.conv33conv33conv11.6.bias', 'layers.0.blocks.1.sk_conv.conv.0.weight', 'layers.0.blocks.1.sk_conv.conv.0.bias', 'layers.0.blocks.1.sk_conv.conv.1.weight', 'layers.0.blocks.1.sk_conv.conv.1.bias', 'layers.0.blocks.1.sk_conv.conv.2.weight', 'layers.0.blocks.1.sk_conv.conv.2.bias', 'layers.0.blocks.1.sk_conv.conv.3.weight', 'layers.0.blocks.1.sk_conv.conv.3.bias', 'layers.0.blocks.1.sk_conv.fc.0.weight', 'layers.0.blocks.1.sk_conv.fc.0.bias', 'layers.0.blocks.1.sk_conv.fc.2.weight', 'layers.0.blocks.1.sk_conv.fc.2.bias', 'layers.1.blocks.0.mutli_head_atten.in_proj_weight', 'layers.1.blocks.0.mutli_head_atten.in_proj_bias', 'layers.1.blocks.0.mutli_head_atten.out_proj.weight', 'layers.1.blocks.0.mutli_head_atten.out_proj.bias', 'layers.1.blocks.0.conv33conv33conv11.0.weight', 'layers.1.blocks.0.conv33conv33conv11.0.bias', 'layers.1.blocks.0.conv33conv33conv11.1.weight', 'layers.1.blocks.0.conv33conv33conv11.1.bias', 'layers.1.blocks.0.conv33conv33conv11.1.running_mean', 'layers.1.blocks.0.conv33conv33conv11.1.running_var', 'layers.1.blocks.0.conv33conv33conv11.3.weight', 'layers.1.blocks.0.conv33conv33conv11.3.bias', 'layers.1.blocks.0.conv33conv33conv11.4.weight', 'layers.1.blocks.0.conv33conv33conv11.4.bias', 'layers.1.blocks.0.conv33conv33conv11.4.running_mean', 'layers.1.blocks.0.conv33conv33conv11.4.running_var', 'layers.1.blocks.0.conv33conv33conv11.6.weight', 'layers.1.blocks.0.conv33conv33conv11.6.bias', 'layers.1.blocks.0.sk_conv.conv.0.weight', 'layers.1.blocks.0.sk_conv.conv.0.bias', 'layers.1.blocks.0.sk_conv.conv.1.weight', 'layers.1.blocks.0.sk_conv.conv.1.bias', 'layers.1.blocks.0.sk_conv.conv.2.weight', 'layers.1.blocks.0.sk_conv.conv.2.bias', 'layers.1.blocks.0.sk_conv.conv.3.weight', 'layers.1.blocks.0.sk_conv.conv.3.bias', 'layers.1.blocks.0.sk_conv.fc.0.weight', 'layers.1.blocks.0.sk_conv.fc.0.bias', 'layers.1.blocks.0.sk_conv.fc.2.weight', 'layers.1.blocks.0.sk_conv.fc.2.bias', 'layers.1.blocks.1.mutli_head_atten.in_proj_weight', 'layers.1.blocks.1.mutli_head_atten.in_proj_bias', 'layers.1.blocks.1.mutli_head_atten.out_proj.weight', 'layers.1.blocks.1.mutli_head_atten.out_proj.bias', 'layers.1.blocks.1.conv33conv33conv11.0.weight', 'layers.1.blocks.1.conv33conv33conv11.0.bias', 'layers.1.blocks.1.conv33conv33conv11.1.weight', 'layers.1.blocks.1.conv33conv33conv11.1.bias', 'layers.1.blocks.1.conv33conv33conv11.1.running_mean', 'layers.1.blocks.1.conv33conv33conv11.1.running_var', 'layers.1.blocks.1.conv33conv33conv11.3.weight', 'layers.1.blocks.1.conv33conv33conv11.3.bias', 'layers.1.blocks.1.conv33conv33conv11.4.weight', 'layers.1.blocks.1.conv33conv33conv11.4.bias', 'layers.1.blocks.1.conv33conv33conv11.4.running_mean', 'layers.1.blocks.1.conv33conv33conv11.4.running_var', 'layers.1.blocks.1.conv33conv33conv11.6.weight', 'layers.1.blocks.1.conv33conv33conv11.6.bias', 'layers.1.blocks.1.sk_conv.conv.0.weight', 'layers.1.blocks.1.sk_conv.conv.0.bias', 'layers.1.blocks.1.sk_conv.conv.1.weight', 'layers.1.blocks.1.sk_conv.conv.1.bias', 'layers.1.blocks.1.sk_conv.conv.2.weight', 'layers.1.blocks.1.sk_conv.conv.2.bias', 'layers.1.blocks.1.sk_conv.conv.3.weight', 'layers.1.blocks.1.sk_conv.conv.3.bias', 'layers.1.blocks.1.sk_conv.fc.0.weight', 'layers.1.blocks.1.sk_conv.fc.0.bias', 'layers.1.blocks.1.sk_conv.fc.2.weight', 'layers.1.blocks.1.sk_conv.fc.2.bias', 'layers.2.blocks.0.mutli_head_atten.in_proj_weight', 'layers.2.blocks.0.mutli_head_atten.in_proj_bias', 'layers.2.blocks.0.mutli_head_atten.out_proj.weight', 'layers.2.blocks.0.mutli_head_atten.out_proj.bias', 'layers.2.blocks.0.conv33conv33conv11.0.weight', 'layers.2.blocks.0.conv33conv33conv11.0.bias', 'layers.2.blocks.0.conv33conv33conv11.1.weight', 'layers.2.blocks.0.conv33conv33conv11.1.bias', 'layers.2.blocks.0.conv33conv33conv11.1.running_mean', 'layers.2.blocks.0.conv33conv33conv11.1.running_var', 'layers.2.blocks.0.conv33conv33conv11.3.weight', 'layers.2.blocks.0.conv33conv33conv11.3.bias', 'layers.2.blocks.0.conv33conv33conv11.4.weight', 'layers.2.blocks.0.conv33conv33conv11.4.bias', 'layers.2.blocks.0.conv33conv33conv11.4.running_mean', 'layers.2.blocks.0.conv33conv33conv11.4.running_var', 'layers.2.blocks.0.conv33conv33conv11.6.weight', 'layers.2.blocks.0.conv33conv33conv11.6.bias', 'layers.2.blocks.0.sk_conv.conv.0.weight', 'layers.2.blocks.0.sk_conv.conv.0.bias', 'layers.2.blocks.0.sk_conv.conv.1.weight', 'layers.2.blocks.0.sk_conv.conv.1.bias', 'layers.2.blocks.0.sk_conv.conv.2.weight', 'layers.2.blocks.0.sk_conv.conv.2.bias', 'layers.2.blocks.0.sk_conv.conv.3.weight', 'layers.2.blocks.0.sk_conv.conv.3.bias', 'layers.2.blocks.0.sk_conv.fc.0.weight', 'layers.2.blocks.0.sk_conv.fc.0.bias', 'layers.2.blocks.0.sk_conv.fc.2.weight', 'layers.2.blocks.0.sk_conv.fc.2.bias', 'layers.2.blocks.1.mutli_head_atten.in_proj_weight', 'layers.2.blocks.1.mutli_head_atten.in_proj_bias', 'layers.2.blocks.1.mutli_head_atten.out_proj.weight', 'layers.2.blocks.1.mutli_head_atten.out_proj.bias', 'layers.2.blocks.1.conv33conv33conv11.0.weight', 'layers.2.blocks.1.conv33conv33conv11.0.bias', 'layers.2.blocks.1.conv33conv33conv11.1.weight', 'layers.2.blocks.1.conv33conv33conv11.1.bias', 'layers.2.blocks.1.conv33conv33conv11.1.running_mean', 'layers.2.blocks.1.conv33conv33conv11.1.running_var', 'layers.2.blocks.1.conv33conv33conv11.3.weight', 'layers.2.blocks.1.conv33conv33conv11.3.bias', 'layers.2.blocks.1.conv33conv33conv11.4.weight', 'layers.2.blocks.1.conv33conv33conv11.4.bias', 'layers.2.blocks.1.conv33conv33conv11.4.running_mean', 'layers.2.blocks.1.conv33conv33conv11.4.running_var', 'layers.2.blocks.1.conv33conv33conv11.6.weight', 'layers.2.blocks.1.conv33conv33conv11.6.bias', 'layers.2.blocks.1.sk_conv.conv.0.weight', 'layers.2.blocks.1.sk_conv.conv.0.bias', 'layers.2.blocks.1.sk_conv.conv.1.weight', 'layers.2.blocks.1.sk_conv.conv.1.bias', 'layers.2.blocks.1.sk_conv.conv.2.weight', 'layers.2.blocks.1.sk_conv.conv.2.bias', 'layers.2.blocks.1.sk_conv.conv.3.weight', 'layers.2.blocks.1.sk_conv.conv.3.bias', 'layers.2.blocks.1.sk_conv.fc.0.weight', 'layers.2.blocks.1.sk_conv.fc.0.bias', 'layers.2.blocks.1.sk_conv.fc.2.weight', 'layers.2.blocks.1.sk_conv.fc.2.bias', 'layers.2.blocks.2.mutli_head_atten.in_proj_weight', 'layers.2.blocks.2.mutli_head_atten.in_proj_bias', 'layers.2.blocks.2.mutli_head_atten.out_proj.weight', 'layers.2.blocks.2.mutli_head_atten.out_proj.bias', 'layers.2.blocks.2.conv33conv33conv11.0.weight', 'layers.2.blocks.2.conv33conv33conv11.0.bias', 'layers.2.blocks.2.conv33conv33conv11.1.weight', 'layers.2.blocks.2.conv33conv33conv11.1.bias', 'layers.2.blocks.2.conv33conv33conv11.1.running_mean', 'layers.2.blocks.2.conv33conv33conv11.1.running_var', 'layers.2.blocks.2.conv33conv33conv11.3.weight', 'layers.2.blocks.2.conv33conv33conv11.3.bias', 'layers.2.blocks.2.conv33conv33conv11.4.weight', 'layers.2.blocks.2.conv33conv33conv11.4.bias', 'layers.2.blocks.2.conv33conv33conv11.4.running_mean', 'layers.2.blocks.2.conv33conv33conv11.4.running_var', 'layers.2.blocks.2.conv33conv33conv11.6.weight', 'layers.2.blocks.2.conv33conv33conv11.6.bias', 'layers.2.blocks.2.sk_conv.conv.0.weight', 'layers.2.blocks.2.sk_conv.conv.0.bias', 'layers.2.blocks.2.sk_conv.conv.1.weight', 'layers.2.blocks.2.sk_conv.conv.1.bias', 'layers.2.blocks.2.sk_conv.conv.2.weight', 'layers.2.blocks.2.sk_conv.conv.2.bias', 'layers.2.blocks.2.sk_conv.conv.3.weight', 'layers.2.blocks.2.sk_conv.conv.3.bias', 'layers.2.blocks.2.sk_conv.fc.0.weight', 'layers.2.blocks.2.sk_conv.fc.0.bias', 'layers.2.blocks.2.sk_conv.fc.2.weight', 'layers.2.blocks.2.sk_conv.fc.2.bias', 'layers.2.blocks.3.mutli_head_atten.in_proj_weight', 'layers.2.blocks.3.mutli_head_atten.in_proj_bias', 'layers.2.blocks.3.mutli_head_atten.out_proj.weight', 'layers.2.blocks.3.mutli_head_atten.out_proj.bias', 'layers.2.blocks.3.conv33conv33conv11.0.weight', 'layers.2.blocks.3.conv33conv33conv11.0.bias', 'layers.2.blocks.3.conv33conv33conv11.1.weight', 'layers.2.blocks.3.conv33conv33conv11.1.bias', 'layers.2.blocks.3.conv33conv33conv11.1.running_mean', 'layers.2.blocks.3.conv33conv33conv11.1.running_var', 'layers.2.blocks.3.conv33conv33conv11.3.weight', 'layers.2.blocks.3.conv33conv33conv11.3.bias', 'layers.2.blocks.3.conv33conv33conv11.4.weight', 'layers.2.blocks.3.conv33conv33conv11.4.bias', 'layers.2.blocks.3.conv33conv33conv11.4.running_mean', 'layers.2.blocks.3.conv33conv33conv11.4.running_var', 'layers.2.blocks.3.conv33conv33conv11.6.weight', 'layers.2.blocks.3.conv33conv33conv11.6.bias', 'layers.2.blocks.3.sk_conv.conv.0.weight', 'layers.2.blocks.3.sk_conv.conv.0.bias', 'layers.2.blocks.3.sk_conv.conv.1.weight', 'layers.2.blocks.3.sk_conv.conv.1.bias', 'layers.2.blocks.3.sk_conv.conv.2.weight', 'layers.2.blocks.3.sk_conv.conv.2.bias', 'layers.2.blocks.3.sk_conv.conv.3.weight', 'layers.2.blocks.3.sk_conv.conv.3.bias', 'layers.2.blocks.3.sk_conv.fc.0.weight', 'layers.2.blocks.3.sk_conv.fc.0.bias', 'layers.2.blocks.3.sk_conv.fc.2.weight', 'layers.2.blocks.3.sk_conv.fc.2.bias', 'layers.2.blocks.4.mutli_head_atten.in_proj_weight', 'layers.2.blocks.4.mutli_head_atten.in_proj_bias', 'layers.2.blocks.4.mutli_head_atten.out_proj.weight', 'layers.2.blocks.4.mutli_head_atten.out_proj.bias', 'layers.2.blocks.4.conv33conv33conv11.0.weight', 'layers.2.blocks.4.conv33conv33conv11.0.bias', 'layers.2.blocks.4.conv33conv33conv11.1.weight', 'layers.2.blocks.4.conv33conv33conv11.1.bias', 'layers.2.blocks.4.conv33conv33conv11.1.running_mean', 'layers.2.blocks.4.conv33conv33conv11.1.running_var', 'layers.2.blocks.4.conv33conv33conv11.3.weight', 'layers.2.blocks.4.conv33conv33conv11.3.bias', 'layers.2.blocks.4.conv33conv33conv11.4.weight', 'layers.2.blocks.4.conv33conv33conv11.4.bias', 'layers.2.blocks.4.conv33conv33conv11.4.running_mean', 'layers.2.blocks.4.conv33conv33conv11.4.running_var', 'layers.2.blocks.4.conv33conv33conv11.6.weight', 'layers.2.blocks.4.conv33conv33conv11.6.bias', 'layers.2.blocks.4.sk_conv.conv.0.weight', 'layers.2.blocks.4.sk_conv.conv.0.bias', 'layers.2.blocks.4.sk_conv.conv.1.weight', 'layers.2.blocks.4.sk_conv.conv.1.bias', 'layers.2.blocks.4.sk_conv.conv.2.weight', 'layers.2.blocks.4.sk_conv.conv.2.bias', 'layers.2.blocks.4.sk_conv.conv.3.weight', 'layers.2.blocks.4.sk_conv.conv.3.bias', 'layers.2.blocks.4.sk_conv.fc.0.weight', 'layers.2.blocks.4.sk_conv.fc.0.bias', 'layers.2.blocks.4.sk_conv.fc.2.weight', 'layers.2.blocks.4.sk_conv.fc.2.bias', 'layers.3.blocks.0.mutli_head_atten.in_proj_weight', 'layers.3.blocks.0.mutli_head_atten.in_proj_bias', 'layers.3.blocks.0.mutli_head_atten.out_proj.weight', 'layers.3.blocks.0.mutli_head_atten.out_proj.bias', 'layers.3.blocks.0.conv33conv33conv11.0.weight', 'layers.3.blocks.0.conv33conv33conv11.0.bias', 'layers.3.blocks.0.conv33conv33conv11.1.weight', 'layers.3.blocks.0.conv33conv33conv11.1.bias', 'layers.3.blocks.0.conv33conv33conv11.1.running_mean', 'layers.3.blocks.0.conv33conv33conv11.1.running_var', 'layers.3.blocks.0.conv33conv33conv11.3.weight', 'layers.3.blocks.0.conv33conv33conv11.3.bias', 'layers.3.blocks.0.conv33conv33conv11.4.weight', 'layers.3.blocks.0.conv33conv33conv11.4.bias', 'layers.3.blocks.0.conv33conv33conv11.4.running_mean', 'layers.3.blocks.0.conv33conv33conv11.4.running_var', 'layers.3.blocks.0.conv33conv33conv11.6.weight', 'layers.3.blocks.0.conv33conv33conv11.6.bias', 'layers.3.blocks.0.sk_conv.conv.0.weight', 'layers.3.blocks.0.sk_conv.conv.0.bias', 'layers.3.blocks.0.sk_conv.conv.1.weight', 'layers.3.blocks.0.sk_conv.conv.1.bias', 'layers.3.blocks.0.sk_conv.conv.2.weight', 'layers.3.blocks.0.sk_conv.conv.2.bias', 'layers.3.blocks.0.sk_conv.conv.3.weight', 'layers.3.blocks.0.sk_conv.conv.3.bias', 'layers.3.blocks.0.sk_conv.fc.0.weight', 'layers.3.blocks.0.sk_conv.fc.0.bias', 'layers.3.blocks.0.sk_conv.fc.2.weight', 'layers.3.blocks.0.sk_conv.fc.2.bias', 'layers.3.blocks.1.mutli_head_atten.in_proj_weight', 'layers.3.blocks.1.mutli_head_atten.in_proj_bias', 'layers.3.blocks.1.mutli_head_atten.out_proj.weight', 'layers.3.blocks.1.mutli_head_atten.out_proj.bias', 'layers.3.blocks.1.conv33conv33conv11.0.weight', 'layers.3.blocks.1.conv33conv33conv11.0.bias', 'layers.3.blocks.1.conv33conv33conv11.1.weight', 'layers.3.blocks.1.conv33conv33conv11.1.bias', 'layers.3.blocks.1.conv33conv33conv11.1.running_mean', 'layers.3.blocks.1.conv33conv33conv11.1.running_var', 'layers.3.blocks.1.conv33conv33conv11.3.weight', 'layers.3.blocks.1.conv33conv33conv11.3.bias', 'layers.3.blocks.1.conv33conv33conv11.4.weight', 'layers.3.blocks.1.conv33conv33conv11.4.bias', 'layers.3.blocks.1.conv33conv33conv11.4.running_mean', 'layers.3.blocks.1.conv33conv33conv11.4.running_var', 'layers.3.blocks.1.conv33conv33conv11.6.weight', 'layers.3.blocks.1.conv33conv33conv11.6.bias', 'layers.3.blocks.1.sk_conv.conv.0.weight', 'layers.3.blocks.1.sk_conv.conv.0.bias', 'layers.3.blocks.1.sk_conv.conv.1.weight', 'layers.3.blocks.1.sk_conv.conv.1.bias', 'layers.3.blocks.1.sk_conv.conv.2.weight', 'layers.3.blocks.1.sk_conv.conv.2.bias', 'layers.3.blocks.1.sk_conv.conv.3.weight', 'layers.3.blocks.1.sk_conv.conv.3.bias', 'layers.3.blocks.1.sk_conv.fc.0.weight', 'layers.3.blocks.1.sk_conv.fc.0.bias', 'layers.3.blocks.1.sk_conv.fc.2.weight', 'layers.3.blocks.1.sk_conv.fc.2.bias'], unexpected_keys=[])\n",
      "loaded successfully\n"
     ]
    }
   ],
   "source": [
    "model = Model(config)\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "model.to(device)\n",
    "model.load_state_dict(torch.load(config.load_model_path))\n",
    "print('loaded successfully')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7f20f670-9076-4fa1-85f3-f946b1b59665",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Left-Diagnostic Keywords:hypertensive retinopathy,Right-Diagnostic Keywords:hypertensive retinopathy\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "----- [Encoding]: 100%|██████████| 1/1 [00:00<00:00, 14.89it/s]\n",
      "----- [Predicting] : 100%|██████████| 1/1 [00:00<00:00,  1.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0 0 0 0 0 1 0 0]]\n",
      "[[0 0 0 0 0 7 2 1]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "from utils.APIs.APIDataset import APIDataset\n",
    "from utils.APIs.APIEncode import api_encode\n",
    "from utils.APIs.APIMetric import api_metric\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "#########################################准备输入文件\n",
    "left_image_path = './28_left_enhance_result.jpg'\n",
    "right_image_path = './28_right_enhance_result.jpg'\n",
    "\n",
    "left_describe = 'hypertensive retinopathy'\n",
    "right_describe = 'hypertensive retinopathy'\n",
    "\n",
    "describe =  \"Left-Diagnostic Keywords:\" + left_describe + \",\" + \"Right-Diagnostic Keywords:\" + right_describe\n",
    "print(describe)\n",
    "label = [0,0,0,0,0,7,2,1]\n",
    "label = np.array(label, dtype=np.float32)\n",
    "label = torch.tensor(label)\n",
    "\n",
    "data = []\n",
    "data.append((left_image_path, right_image_path, describe, label))\n",
    "\n",
    "#########################################转化为Dataloader\n",
    "dataset_inputs = api_encode(data,config,'test')\n",
    "dataset = APIDataset(*dataset_inputs)\n",
    "dataloader = DataLoader(dataset=dataset, **config.checkout_params, collate_fn=dataset.collate_fn, drop_last=False)\n",
    "\n",
    "#########################################预测\n",
    "model.eval()\n",
    "pred_labels = None\n",
    "true_labels = None\n",
    "threshold = 0.5\n",
    "for batch in tqdm(dataloader, desc='----- [Predicting] '):\n",
    "    texts, texts_mask, imgs, labels = batch\n",
    "    texts, texts_mask, imgs = texts.to(device), texts_mask.to(device), imgs.to(device)\n",
    "    pred_vec = model(texts, texts_mask, imgs)\n",
    "\n",
    "    pred_labels = pred_vec.cpu().detach().numpy()\n",
    "    true_labels = labels.cpu().detach().numpy()\n",
    "\n",
    "    predicted_classes = (pred_labels >= threshold).astype(int)\n",
    "    print(predicted_classes)\n",
    "    print(true_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e1abb2-2179-410d-bad4-3c906d0c842d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a52e37-cb84-4de8-86cb-381c131d9152",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff430f27-42fe-41cb-91dc-150d71f582bf",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21d6b8d0-d80e-4911-a55d-7714daf49d2d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13c0acd0-b5ca-4a63-afbd-6706f7765b5c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91478206-18c9-45f0-a4f2-893b13c1f3dd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
